import torch
from torchvision.models.resnet import resnet18
from collections import defaultdict
from torch.utils.tensorboard._pytorch_graph import NodePy, NodePyIO, NodePyOP, GraphPy
import re

dummpy_input = torch.randn((8, 3, 224, 224)).cuda()
model = resnet18(True).eval().cuda()

# 生成torchscript
if torch.__version__ >= "1.6.0":
    trace = torch.jit.trace(model, dummpy_input, strict=False)
else:
    trace = torch.jit.trace(model, dummpy_input)

graph = trace.graph      # 该变量为引用类型，tace.graph改变也会改变graph的值
# print(graph)
# print("*"*40)
# 将submodule展开
torch._C._jit_pass_inline(trace.graph)

# print(graph)
# print("*"*40)

# 映射输入 输出到节点
output_to_node = defaultdict(list)   # 表示某个输出是某个节点的输出
input_to_node = defaultdict(list)    # 表示某个输入时某个节点的输入

for node in graph.nodes():
    # print(type(node))
    # 如果节点类型是常量，则跳过
    if node.kind() == 'prim::Constant':
        continue
    # 遍历当前节点的输入，并保存
    for x in node.inputs():
        # x的类型是torch._C.Value
        # print(type(x))
        # print(node)
        # print(x.node())
        # x.node()表示该value来自哪个节点
        if x.node().kind() == 'prim::Constant':
            continue
        input_to_node[x.debugName()].append(node)
    # 遍历当前节点的输出，并保存，一个输出只能来自一个节点
    for x in node.outputs():
        if x.node().kind() == 'prim::Constant':
            continue
        output_to_node[x.debugName()].append(node)
        assert len(output_to_node[x.debugName()]) <= 1, "一个输出只能来自一个节点!"


nodes_py = GraphPy()
# 获得graph的输入和输出节点
for node in graph.inputs():
    # print(node)
    # print(node.type())
    # print(node.type().kind())
    # print("*"*40)
    # 过滤没有使用的节点
    if not node.uses():
        continue

    # 过滤网络整体节点resnet，该值一般有TensorType和ClassType等几种
    if node.type().kind() != "ClassType":
        nodes_py.append(NodePyIO(node, "input"))

# for node in graph.outputs():
#     print(node)
#     print(node.type())
#     print(node.type().kind())
#     print("*"*40)


def extract_leaf_modules(trace):
    # 获得叶子模块，不存在子模块
    def is_parent(name1, name2):
        """
        check if name1 is parent node of name2, for example:
        name1: aa.bb,  name2: aa.bb.cc,  return True
        name1: aa.b,  name2: aa.bb, return False
        """
        parts1, parts2 = name1.split('.'), name2.split('.')
        if len(parts1) >= len(parts2):
            return False
        for i, _ in enumerate(parts1):
            if parts2[i] != parts1[i]:
                return False
        return True

    module_names = sorted([x[0] for x in trace.named_modules() if x[0]])
    leaf_nodes = []
    for i, name in enumerate(module_names):
        if i + 1 >= len(module_names) or not is_parent(name, module_names[i + 1]):
            leaf_nodes.append(name)
    return leaf_nodes

# trace_modules = {}
# model_modules = {}
# for name, module in trace.named_modules():
#     # print(x)
#     trace_modules[name] = module
#     # print("x"*40)

# for name, module in model.named_modules():
#     model_modules[name] = module

# for name in trace_modules.keys():
#     print(name)
#     print(trace_modules[name])
#     print(model_modules[name])
#     print("x"*40)


leaf_modules = extract_leaf_modules(trace)
print(leaf_modules)
print("x"*40)

def parse_traced_name(module_name):
    prefix = 'TracedModule['
    suffix = ']'
    if module_name.startswith(prefix) and module_name.endswith(suffix):
        module_name = module_name[len(prefix):-len(suffix)]
    return module_name

# 统计module对应的类别(如conv bn relu等)
module_to_type = {name: parse_traced_name( module._name) for name, module in trace.named_modules()}


def get_module_name(scope_name):
    """
    根据节点的scopeName检索对应module的name
    Parameters:
    -----------
    scope_name: str
        scope_name of a graph node, for example:
        for pytorch 1.3.1: MyModel/BackboneModel[backbone]/Conv2d[conv2]
        for pytorch 1.4.0: __module.backbone/__module.backbone.conv2

    Returns:
    -------
    str
        module name, such as backbone.conv2
    """
    if torch.__version__ >= '1.4.0':
        return scope_name.split('/')[-1].replace('__module.', '')
    else:
        return '.'.join(re.findall(r'\[(.*?)\]', scope_name))


# 关联module或者func和其对应的node
# 关联module与其对应的所有节点
module_to_nodes = defaultdict(list)
# 关联函数与其对应的所有节点
func_to_nodes = defaultdict(list)

for node in graph.nodes():
    if node.kind() == 'prim::Constant':
        continue
    scopeName = node.scopeName()
    moduleName = get_module_name(scopeName)
    print(scopeName, "对应的module或者func的Name为：", moduleName)
    if moduleName in leaf_modules:
        module_to_nodes[moduleName].append(node)
    else:
        func_to_nodes[scopeName].append(node)

